import functools
import timeit
import memory_profiler
import datetime

from loguru import logger


def print_execution_time(date_format):
    def decorator(func):
        def wrapper(*args, **kwargs):
            start_time = datetime.datetime.now()
            start_time_str = start_time.strftime(date_format)
            logger.info(f"Function '{func.__name__}' started at: {start_time_str}")

            result = func(*args, **kwargs)

            end_time = datetime.datetime.now()
            end_time_str = end_time.strftime(date_format)
            logger.info(f"Function '{func.__name__}' completed at: {end_time_str}")

            return result

        return wrapper

    return decorator


def calculate_execution_time(func):
    # 通过装饰器，返回函数执行时间
    @functools.wraps(func)  # 保留被装饰函数的元信息
    def wrapper(*args, **kwargs):
        start_time = timeit.default_timer()

        result = func(*args, **kwargs)

        end_time = timeit.default_timer()
        execution_time = end_time - start_time
        logger.info(f"Function '{func.__name__}' executed in {execution_time} seconds")

        return result

    return wrapper


def calculate_execution_memory(func):
    # 通过装饰器，返回函数执行过程中占用的内存
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        mem_profile = memory_profiler.memory_usage()
        start_mem = max(mem_profile)

        result = func(*args, **kwargs)

        mem_profile = memory_profiler.memory_usage()
        end_mem = max(mem_profile)
        avg_mem = (start_mem + end_mem) / 2

        logger.info(f"Function '{func.__name__}' memory usage: Peak = {end_mem} MiB, Average = {avg_mem} MiB")

        return result

    return wrapper
